package io.gitee.welson.sensitive.writer;

import com.alibaba.excel.EasyExcel;
import com.alibaba.excel.ExcelWriter;
import com.alibaba.excel.converters.Converter;
import com.alibaba.excel.write.metadata.WriteSheet;
import io.gitee.welson.sensitive.config.SensitiveExportProperties;
import io.gitee.welson.sensitive.converter.StringSensitiveConverter;

import java.io.OutputStream;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 默认的Excel导出写入器实现
 * 基于EasyExcel实现，支持分批写入、自动分页和多sheet页写入
 *
 * @author welson
 */
public class DefaultExportWriter implements ExportWriter {

    private ExcelWriter excelWriter;
    private final Converter<String> converter;
    private final Map<String, WriteSheet> sheetMap;
    private final Map<String, Integer> sheetIndexMap;
    private final Map<String, AtomicInteger> rowCountMap;
    private int maxRowsPerSheet;
    private final int defaultMaxRows;
    private Class<?> headClass;
    private static final String DEFAULT_SHEET_PREFIX = "Sheet";

    public DefaultExportWriter(SensitiveExportProperties properties, StringSensitiveConverter converter) {
        this.sheetMap = new ConcurrentHashMap<>();
        sheetIndexMap = new ConcurrentHashMap<>();
        this.rowCountMap = new ConcurrentHashMap<>();
        this.defaultMaxRows = properties.getSheetMaxRows();
        this.maxRowsPerSheet = this.defaultMaxRows;
        this.converter = converter;
    }

    @Override
    public void setSheetMaxRows(int maxRows) {
        if (maxRows <= 0) {
            // 如果设置的值无效，使用默认值
            this.maxRowsPerSheet = this.defaultMaxRows;
        } else {
            this.maxRowsPerSheet = maxRows;
        }
    }

    @Override
    public <T> void initialize(OutputStream outputStream, Class<T> dataClass) {
        this.headClass = dataClass;
        this.excelWriter = EasyExcel.write(outputStream, dataClass).registerConverter(converter).build();
        // 重置为默认值，确保每次导出任务都从默认值开始
        this.maxRowsPerSheet = this.defaultMaxRows;
        // 清理上次导出的状态
        this.sheetMap.clear();
        this.rowCountMap.clear();
    }

    @Override
    public void writeBatch(List<?> dataList) {
        writeToSheet(DEFAULT_SHEET_PREFIX + "1", dataList);
    }

    @Override
    public void writeBatchWithPaging(List<?> dataList) {
        writeToSheetWithPaging(DEFAULT_SHEET_PREFIX + "1", dataList);
    }

    @Override
    public void writeToSheet(String sheetName, List<?> dataList) {
        WriteSheet writeSheet = getOrCreateSheet(sheetName);
        excelWriter.write(dataList, writeSheet);
    }

    @Override
    public void writeToSheetWithPaging(String sheetName, List<?> dataList) {
        if (dataList == null || dataList.isEmpty()) {
            return;
        }

        // 递归处理数据，确保每个sheet不超过最大行数
        doWriteToSheetWithPaging(sheetName, dataList, 0);
    }

    private void doWriteToSheetWithPaging(String sheetName, List<?> dataList, int baseIndex) {
        WriteSheet writeSheet = getOrCreateSheet(sheetName, baseIndex);
        AtomicInteger rowCount = rowCountMap.get(writeSheet.getSheetName());

        // 如果Sheet页加上当前批次数据后超过单个sheet最大行数
        if (rowCount.get() + dataList.size() > maxRowsPerSheet) {
            // 写入当前sheet最大容量的数据
            List<?> currentBatch = dataList.subList(0, maxRowsPerSheet - rowCount.get());
            excelWriter.write(currentBatch, writeSheet);
            rowCount.set(maxRowsPerSheet);

            // 递归处理剩余数据
            List<?> remainingData = dataList.subList(maxRowsPerSheet - rowCount.get(), dataList.size());
            if (!remainingData.isEmpty()) {
                if (baseIndex == 0) {
                    writeSheet.setSheetName(sheetName.replaceAll("\\d+$", "") + 1);
                }
                doWriteToSheetWithPaging(sheetName, remainingData, baseIndex + 1);
            }
        } else {
            // 当前批次数据未超过单个sheet最大行数
            excelWriter.write(dataList, writeSheet);
            rowCount.addAndGet(dataList.size());
        }
    }

    private WriteSheet getOrCreateSheet(String sheetName) {
        Integer index = sheetIndexMap.computeIfAbsent(sheetName, k -> sheetIndexMap.size());
        return sheetMap.computeIfAbsent(sheetName,
                k -> EasyExcel.writerSheet(index, sheetName)
                .head(headClass)
                .build());
    }

    private WriteSheet getOrCreateSheet(String baseSheetName, int index) {
        String actualSheetName = index == 0 ? baseSheetName : baseSheetName.replaceAll("\\d+$", "") + (index + 1);
        
        return sheetMap.computeIfAbsent(actualSheetName, k -> {
            WriteSheet writeSheet = EasyExcel.writerSheet(index, k)
                    .head(headClass)
                    .build();
            rowCountMap.put(k, new AtomicInteger(0));
            return writeSheet;
        });
    }

    @Override
    public void finish() {
        if (excelWriter != null) {
            excelWriter.finish();
        }
    }

    @Override
    public void close() {
        finish();
    }

    @Override
    public void flush() {
        // EasyExcel handles buffering internally
    }
} 